# Copyright (c) 2008, RPI-ACM, Jesse Kempf <jessekempf@gmail.com>
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions
# are met:
# - Redistributions of source code must retain the above copyright notice,
#     this list of conditions and the following disclaimer.
# - Redistributions in binary form must reproduce the above copyright
#     notice, this list of conditions and the following disclaimer in the
#     documentation and/or other materials provided with the distribution.
# - Neither the names of RPI-ACM nor the names of its contributors
#     may be used to endorse or promote products derived from this software
#     without specific prior written permission.
#
#  THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
#  "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED
#  TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A
#  PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER
#  OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
#  EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
#  PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
#  PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
#  LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING
#  NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
#  SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

"""A collection of functions to make it easier to write
privilege-separated unix daemons."""

# pylint: disable-msg=W0142

__all__ = ['daemon', 'revoke_privs', 'privsep', 'fork_child',
			'ProcessContext', 'PIDFile', 'MAXFD']

import grp
import os
import pwd
import signal
import sys
import resource
import errno
import fcntl

import unixutils.setproctitle

from unixutils.privsep.command import CommandPackage
import unixutils.privsep.cmdchannel as cmdchannel

DEFAULT_CHROOT = "/var/empty"

# MAXFD is the soft limit on file descriptors at the time this module
# is loaded.

(MAXFD, MAXFD_HARD) = resource.getrlimit(resource.RLIMIT_NOFILE)

class PrivilegeRevocationFailure(OSError):
	"""It was not possible to revoke privileges for some reason."""
	pass

class PIDFileError(IOError):
	"""It was not possible to acquire a PID file."""
	pass

def _current_uname():
	"""Get the user name associated with the current process's uid."""
	return pwd.getpwuid(os.getuid()).pw_name

def _current_gname():
	"""Get the group name associated with the current process's gid."""
	return grp.getgrgid(os.getgid()).gr_name
	
# pylint: disable-msg=R0903,R0913
class ProcessContext(object):
	"""
		A wrapper for a self-contained process.

		Two functions are associated with the process context. In terms of run
		order, the first is a prep function to be run before the parent forks.
		The second becomes the main function for the child process.
	"""
	def __init__(self, main, pre=None, args=None, kwargs=None, debug=False):
		"""
		Create the ProcessContext.

		f() is callable and will be the main function of the resulting
		  process. Its first argument will always be provided as a
		  PyNIX privsep CommandChannel.
		pre() is an optional callable (with no arguments) and will be
		  run in the parent before it forks to start f().
		args and kwargs are automatically unpacked and passed to f when
		  it is called.
		"""
		self.main = main
		self.pre = pre if pre else (lambda: 1)
		self.args = args if args is not None else []
		self.kwargs = kwargs if kwargs is not None else {}
		self.debug = debug

	def __call__(self, channel):
		self.main(channel, *self.args, **self.kwargs)
# pylint: enable-msg=R0903,R0913

def redirect_stdio():
	"""Redirect stdin, stdout, and stderr to /dev/null."""

	devnull = os.fdopen(os.open(os.devnull, os.O_RDWR), 'rw')
	for stream in (sys.stdin, sys.stdout, sys.stderr):
		os.dup2(stream.fileno(), devnull.fileno())
	sys.stdin = devnull
	sys.stdout = devnull
	sys.stderr = devnull

def close_fds(preserve):
	"""Close all file descriptors not in the preserve list."""

	for fdesc in range(sys.stderr.fileno() + 1, MAXFD + 1):
		if fdesc not in preserve:
			try:
				os.close(fdesc)
			except OSError:
				pass

def daemon(preserve=None, rdr_stdio=True, newcwd='/'):
	"""
		Detach from the controlling terminal and run in the background.

		rdr_stdio controls whether stdin, stdout, and stderr are
		redirected to /dev/null. The daemon's new working directory is
		controlled by newcwd.

		The preserve argument is a list of non-stdio descriptors to skip
		closing. If preserve is set to the special value 'ALL' then
		daemon() will skip closing all remaining fds.
	"""

	preserve = preserve if preserve is not None else []

	# Temporarily ignore SIGHUP
	osh = signal.getsignal(signal.SIGHUP)
	signal.signal(signal.SIGHUP, signal.SIG_IGN)

	# We fork, become a session leader, and fork again.
	# We use os._exit() to avoid strange effects from
	# exception handlers being called.

	# Die if parent
	# pylint: disable-msg=W0212
	if os.fork() != 0:
		os._exit(0)

	os.setsid() # Become session leader

	# Die if parent, again.
	if os.fork() != 0:
		os._exit(0)
	# pylint: enable-msg=W0212

	# Stop ignoring SIGHUP
	signal.signal(signal.SIGHUP, osh)

	# Switch somewhere (to the root of the filesystem by default)
	# to avoid keeping any unneeded directories in use.

	os.chdir(newcwd)

	# Duplicate /dev/null to the stdio fds.
	if rdr_stdio:
		redirect_stdio()

	if not preserve == 'ALL':
		close_fds(preserve)

# pylint: disable-msg=R0913
def privsep(master, child, commands, user, group, root, rdr_master_stdio=True):
	"""Create a privilege separated daemon.

	Both master and child must be provided ProcessContexts. The next
	three arguments control the environment in which the child runs, and
	the last controls whether the master process closes its open FDs.
	"""

	daemon(rdr_stdio=rdr_master_stdio)

	(cchan, _) = fork_child(child, commands, user, group, root, child.debug)
	master.pre()
	master(cchan)

def fork_child(child, commands, user, group, root, debug=False):
	"""Fork off a privilege-revoking child process and return a tuple of
	(commchan, pid).

	A ProcessContext for the child must be provided, along with a list of
	Commands or Messages to be loaded into the CommandChannel that will be
	opened between the master and child. The next three parameters control
	the user, group, and root for the child process.

	Setting the user or group to the special value '-' will leave the
	child running under the original user or group id respectively.

	Setting the debug option will force the child to ignore SIGINT (which
	is delivered to all children of a process that is ^C-ed) and prevents
	file descriptor closing and stdio redirection.	
	"""

	cchan = cmdchannel.CommandChannel()
	cchan.load_commands(CommandPackage(*commands))

	if user == '-':
		user = _current_uname()
	if group == '-':
		group = _current_gname()

	fork = os.fork()
	if fork:
		# If we're the parent, we're the master.
		# We tell the command channel that and
		# return it to the caller.
		cchan.set_parent()
		return (cchan, fork)

	else:
		# We're the child. We tell the command channel that.
		# We then close all unneeded FDs and invoke the
		# pre function which should handle imports of any particular
		# any particular modules the child needs, then close all
		# unneeded FDs, just to be sure.
		# If we're in debug mode, we block SIGINT from ^C-ing the
		# parent process.
		# Then we chroot, drop privs, and then enter the child function.
		cchan.set_child()

		if not debug:
			redirect_stdio()
			close_fds(preserve=[cchan.fileno()])

		child.pre()

		if not debug:
			close_fds(preserve=[cchan.fileno()])

		os.chdir(root)
		revoke_privs(user, group, root)

		if debug:
			signal.signal(signal.SIGINT, signal.SIG_IGN)

		child(cchan)
		os._exit(0)

# pylint: enable-msg=R0913

def revoke_privs(user, group, chroot=DEFAULT_CHROOT):
	"""
	Revoke root privileges of the running process.

	Set the real, effective, and saved set-user-ID and the real, effective,
	and saved set-group-ID to the specified values after chrooting into the
	specified directory.

	This must, of course, be run as root.

	Let uid be the UID of the given user. Let gid be the GID of the
	given group.

	revoke_privs() guarantees that either:
		process uid == process euid == uid
		and
		process gid == process egid == gid
	or
		a PrivilegeRevocationFailure will be raised.
	"""

	pwent = pwd.getpwnam(user)
	grent = grp.getgrnam(group)

	uid = pwent.pw_uid
	gid = grent.gr_gid

	os.chdir(chroot)
	os.chroot(chroot)

	os.setgid(gid)
	os.setuid(uid)

	if not os.getuid() == uid:
		raise PrivilegeRevocationFailure(
			"Changing uid failed (still %d)" % uid
		)
	if not os.geteuid() == uid:
		raise PrivilegeRevocationFailure(
			"Changing euid failed (still %d)" % uid
		)
	if not os.getgid() == gid:
		raise PrivilegeRevocationFailure(
			"Changing gid failed (still %d)" % gid
		)
	if not os.getegid() == gid:
		raise PrivilegeRevocationFailure(
			"Changing egid failed (still %d)" % gid
		)

def setproctitle(title):
	"""
	Set process title.

	Uses underlying system facilities to change the process title that appears
	in the ps(1) command. For more information, see the manpage
	for setproctitle(3).
	"""

	unixutils.setproctitle.setproctitle(title)

def getproctitle():
	"""
	Obtain the process title.
	"""

	return unixutils.setproctitle.getproctitle()

class PIDFile(object):
	"""A context manager representing Unix PID files.

	PID files are used for the following purposes in Unix:
	* Providing a lock so more than one daemon doesn't run at the
		same time.
	* Providing a mechanism for init / rc.d / SMF / &c to determine
		the process IDs of daemons they've started.

	daemon()

	pidfile = PIDFile("/var/run/mydaemon.pid")

	try:
		pidfile.acquire() # acquire the PID file or raise PIDFileError.
	except PIDFileError:
		sys.exit(1) # abort execution

	with pidfile:
		revoke_privs(...)
		...	# do daemon stuff
	... # other cleanup

	PIDFile follows the ContextManager protocol so that abnormal program
	terminations from unchecked exceptions don't leave stale PID files
	laying around. NB: This guarantee can not protect against stale PID
	files as a result of interpreter crashes or calls to os._exit() (both
	of which bypass proper interpreter shutdown).

	To work correctly, the PIDFile must be created after daemon()izing, but
	before a call to revoke_privs().
	"""

	def __init__(self, filename):
		"""Create a PIDFile using the provided file name."""

		self.pidfd = None
		self.filename = filename

	def acquire(self):
		"""Open, lock, and write out a PID file.

		Raises PIDFileException if another process holds the PID file
		instead."""

		self.pidfd = os.open(self.filename, os.O_RDWR | os.O_CREAT)
		try:
			fcntl.lockf(self.pidfd, fcntl.LOCK_EX | fcntl.LOCK_NB)
		except IOError, exc:
			if exc.errno in (errno.EACCES, errno.EAGAIN):
				raise PIDFileError(exc)
			else:
				raise

		os.ftruncate(self.pidfd, 0)
		os.write(self.pidfd, '%d\n' % os.getpid())
		os.fsync(self.pidfd)

	def __enter__(self):
		return self

	def __exit__(self, exc_type, exc_val, exc_tb):
		self.release()

	def release(self):
		"""Release the PID file and delete it."""

		os.unlink(self.filename)
		os.close(self.pidfd)
